package main

import (
	"encoding/json"
	"fmt"
	"github.com/google/uuid"
	"github.com/gorilla/websocket"
	"github.com/pion/rtcp"
	"github.com/pion/webrtc/v4"
	"log"
	"net"
	"net/http"
	"sync"
	"time"
)

type Signal struct {
	Type      string                   `json:"type"`
	SDP       string                   `json:"sdp,omitempty"`
	Candidate *webrtc.ICECandidateInit `json:"candidate,omitempty"`
}

type Room struct {
	PublisherPC     *webrtc.PeerConnection
	PublisherTracks map[string]*webrtc.TrackLocalStaticRTP   // trackID -> local track (fanout source)
	TrackFanOut     map[string][]*webrtc.TrackLocalStaticRTP // trackID -> subscriber tracks
	RemoteTracks    map[string]*webrtc.TrackRemote           // trackID -> remote track
	Subscribers     map[string]*webrtc.PeerConnection        // subscriberID -> PC
	mu              sync.Mutex
}

var (
	rooms    = make(map[string]*Room)
	roomsMu  sync.Mutex
	upgrader = websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool { return true },
	}
)

type SafeConn struct {
	Conn *websocket.Conn
	mu   sync.Mutex
}

func (s *SafeConn) WriteJSON(v interface{}) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.Conn.WriteJSON(v)
}

func main() {
	http.Handle("/", http.FileServer(http.Dir("static")))
	http.HandleFunc("/ws/publish", handlePublish)
	http.HandleFunc("/ws/subscribe", handleSubscribe)
	fmt.Println("🚀 WebRTC SFU Server started at :8180")
	log.Fatal(http.ListenAndServe(":8180", nil))
}

func getOrCreateRoom(roomID string) *Room {
	roomsMu.Lock()
	defer roomsMu.Unlock()
	if room, ok := rooms[roomID]; ok {
		return room
	}
	room := &Room{
		PublisherTracks: make(map[string]*webrtc.TrackLocalStaticRTP),
		TrackFanOut:     make(map[string][]*webrtc.TrackLocalStaticRTP),
		RemoteTracks:    make(map[string]*webrtc.TrackRemote),
		Subscribers:     make(map[string]*webrtc.PeerConnection),
	}
	rooms[roomID] = room
	return room
}

func handlePublish(w http.ResponseWriter, r *http.Request) {
	roomID := r.URL.Query().Get("room")
	if roomID == "" {
		http.Error(w, "Missing room ID", http.StatusBadRequest)
		return
	}
	room := getOrCreateRoom(roomID)

	rawConn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("❌ WebSocket upgrade error:", err)
		return
	}
	scone := &SafeConn{Conn: rawConn}

	pc, err := webrtc.NewPeerConnection(webrtc.Configuration{})
	if err != nil {
		log.Println("Failed to create PeerConnection:", err)
		return
	}

	pc.OnICECandidate(func(c *webrtc.ICECandidate) {
		if c != nil {
			js := c.ToJSON()
			_ = scone.WriteJSON(Signal{Type: "candidate", Candidate: &js})
		}
	})

	pc.OnTrack(func(remoteTrack *webrtc.TrackRemote, receiver *webrtc.RTPReceiver) {
		trackID := remoteTrack.ID()
		localTrack, err := webrtc.NewTrackLocalStaticRTP(remoteTrack.Codec().RTPCodecCapability, trackID, "pion")
		if err != nil {
			log.Println("Create local track error:", err)
			return
		}

		room.mu.Lock()
		room.PublisherTracks[trackID] = localTrack
		room.RemoteTracks[trackID] = remoteTrack
		room.TrackFanOut[trackID] = []*webrtc.TrackLocalStaticRTP{localTrack} // seed
		room.mu.Unlock()

		go forwardRTP(roomID, trackID)
		go sendPLI(room, remoteTrack)
	})

	room.mu.Lock()
	room.PublisherPC = pc
	room.mu.Unlock()

	for {
		_, msg, err := scone.Conn.ReadMessage()
		if err != nil {
			break
		}
		var sig Signal
		if err := json.Unmarshal(msg, &sig); err != nil {
			continue
		}
		switch sig.Type {
		case "offer":
			_ = pc.SetRemoteDescription(webrtc.SessionDescription{Type: webrtc.SDPTypeOffer, SDP: sig.SDP})
			answer, _ := pc.CreateAnswer(nil)
			_ = pc.SetLocalDescription(answer)
			_ = scone.WriteJSON(Signal{Type: "answer", SDP: answer.SDP})
		case "candidate":
			if sig.Candidate != nil {
				_ = pc.AddICECandidate(*sig.Candidate)
			}
		}
	}
}

func handleSubscribe(w http.ResponseWriter, r *http.Request) {
	roomID := r.URL.Query().Get("room")
	if roomID == "" {
		roomID = "default"
	}
	room := getOrCreateRoom(roomID)

	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("WebSocket upgrade error:", err)
		return
	}
	scone := &SafeConn{Conn: ws}
	subscriberID := uuid.NewString()

	peerConn, err := webrtc.NewPeerConnection(webrtc.Configuration{})
	if err != nil {
		log.Println("❌ Create PeerConnection error:", err)
		return
	}

	peerConn.OnICECandidate(func(c *webrtc.ICECandidate) {
		if c != nil {
			js := c.ToJSON()
			_ = scone.WriteJSON(Signal{Type: "candidate", Candidate: &js})
		}
	})

	room.mu.Lock()
	for trackID, src := range room.PublisherTracks {
		capCodec := webrtc.RTPCodecCapability{
			MimeType:     src.Codec().MimeType,
			ClockRate:    src.Codec().ClockRate,
			Channels:     src.Codec().Channels,
			SDPFmtpLine:  src.Codec().SDPFmtpLine,
			RTCPFeedback: src.Codec().RTCPFeedback,
		}
		subscriberTrack, _ := webrtc.NewTrackLocalStaticRTP(capCodec, trackID, "pion")
		sender, _ := peerConn.AddTrack(subscriberTrack)
		go readRTCP(sender)
		room.TrackFanOut[trackID] = append(room.TrackFanOut[trackID], subscriberTrack)
	}
	room.Subscribers[subscriberID] = peerConn
	room.mu.Unlock()

	offer, _ := peerConn.CreateOffer(nil)
	_ = peerConn.SetLocalDescription(offer)
	<-webrtc.GatheringCompletePromise(peerConn)
	_ = scone.WriteJSON(Signal{Type: "offer", SDP: peerConn.LocalDescription().SDP})

	//PublisherChannel, err := room.PublisherPC.CreateDataChannel("mataInfo", nil)
	//PublisherChannel.OnOpen(func() {
	//
	//})
	if err != nil {
		return
	}

	for {
		_, msgBytes, err := ws.ReadMessage()
		if err != nil {
			break
		}
		var sig Signal
		if err := json.Unmarshal(msgBytes, &sig); err != nil {
			continue
		}
		switch sig.Type {
		case "answer":
			_ = peerConn.SetRemoteDescription(webrtc.SessionDescription{Type: webrtc.SDPTypeAnswer, SDP: sig.SDP})
		case "candidate":
			if sig.Candidate != nil {
				_ = peerConn.AddICECandidate(*sig.Candidate)
			}
		}
	}

	room.mu.Lock()
	delete(room.Subscribers, subscriberID)
	room.mu.Unlock()
}

func forwardRTP(roomID, trackID string) {
	room := getOrCreateRoom(roomID)
	buf := make([]byte, 1500)
	for {
		n, _, err := room.RemoteTracks[trackID].Read(buf)
		if err != nil {
			return
		}
		room.mu.Lock()
		for _, t := range room.TrackFanOut[trackID] {
			_, _ = t.Write(buf[:n])
		}
		room.mu.Unlock()
	}
}

func readRTCP(sender *webrtc.RTPSender) {
	rtcpBuf := make([]byte, 1500)
	for {
		_, _, err := sender.Read(rtcpBuf)
		if err != nil {
			return
		}
	}
}

func sendPLI(room *Room, track *webrtc.TrackRemote) {
	ticker := time.NewTicker(time.Second * 3)
	defer ticker.Stop()
	for range ticker.C {
		room.mu.Lock()
		if room.PublisherPC != nil {
			_ = room.PublisherPC.WriteRTCP([]rtcp.Packet{
				&rtcp.PictureLossIndication{MediaSSRC: uint32(track.SSRC())},
			})
		}
		room.mu.Unlock()
	}
}

func startRTSPPush(track *webrtc.TrackRemote) {
	addr, err := net.ResolveUDPAddr("udp", "rtsp://192.168.2.16:8554/mystream") // 推送到 RTSP Simple Server 的视频端口
	if err != nil {
		log.Println("resolve UDP err:", err)
		return
	}
	conn, err := net.DialUDP("udp", nil, addr)
	if err != nil {
		log.Println("dial UDP err:", err)
		return
	}

	log.Println("📡 RTSP Push started...")

	buf := make([]byte, 1500)
	for {
		n, _, readErr := track.Read(buf)
		if readErr != nil {
			log.Println("track read err:", readErr)
			return
		}
		pkt := make([]byte, n)
		copy(pkt, buf[:n])
		_, err := conn.Write(pkt)
		if err != nil {
			log.Println("rtp write err:", err)
			return
		}
	}
}
